package frege.data.HashMap where

import Data.List (lookup)
import Data.Bits (.&., shiftR)

data HashMap a b = Empty 
                 | ! Branch { b0, b1, b2, b3, b4, b5, b6, b7, b8, b9, bA, bB, bC, bD, bE, bF :: HashMap a b }
                 | ! Bucket [(a,b)]
               
derive Eq HashMap a b  

instance Show (Show a, Show b) => HashMap a b where
    show s = "fromList " ++ show (entries s)            

instance Empty HashMap a where   
  empty = Empty
  null Empty = true
  null _ = false

singleton :: Eq a => a -> b -> HashMap a b
singleton k v = insert k v empty

insert :: Eq a => a -> b -> HashMap a b -> HashMap a b
insert k v s = ins (hashCode k) 0 s where
   ins _ 7 Empty = Bucket [(k,v)]
   ins _ 7 (bucket @ Bucket xs) = Bucket ((k,v) : filter ((!=k) <~ fst) xs)
   ins h k Empty = changeBranch emptyBranch (h `.&.` 15) (const $ ins (h `shiftR` 4) (k + 1) Empty)
   ins h k branch = changeBranch branch (h `.&.` 15) $ ins (h `shiftR` 4) (k + 1)

delete :: Eq a => a -> HashMap a b -> HashMap a b
delete k s = del (hashCode k) s where
   del _ Empty = Empty
   del _ (Bucket xs) = emptyOrBucket $ filter ((!=k) <~ fst) xs
   del h branch = replaceEmptyBranch $ changeBranch branch (h `.&.` 15) $ del (h `shiftR` 4)
   replaceEmptyBranch (Branch Empty Empty Empty Empty Empty Empty Empty Empty
                              Empty Empty Empty Empty Empty Empty Empty Empty) = Empty
   replaceEmptyBranch branch = branch
       
contains :: Eq a => a -> HashMap a b -> Bool
contains k s = cnt s (hashCode k) where
   cnt Empty _ = false
   cnt (Bucket xs) _ = elem k $ map fst xs  
   cnt branch h = cnt (getBranch branch (h `.&.` 15)) (h `shiftR` 4)
      
get :: Eq a => a -> HashMap a b -> Maybe b
get k s = get' s (hashCode k) where
   get' Empty _ = Nothing
   get' (Bucket xs) _ = lookup k xs  
   get' branch h = get' (getBranch branch (h `.&.` 15)) (h `shiftR` 4)
      
partition :: Eq a => (a -> b -> Bool) -> HashMap a b -> (HashMap a b, HashMap a b)
partition f s = foldr insPair (empty, empty) $ entries s where
    insPair (k,v) (left, right) | f k v = (insert k v left, right)
                                | otherwise = (left, insert k v right)   
                                
instance Functor Eq k => HashMap k where
   fmap _ Empty = Empty
   fmap f (Bucket xs) = Bucket $ map (\(k,v) -> (k,f v)) xs 
   fmap f (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF) = 
         Branch (fmap f x0) (fmap f x1) (fmap f x2) (fmap f x3) (fmap f x4) (fmap f x5) (fmap f x6) (fmap f x7) 
                (fmap f x8) (fmap f x9) (fmap f xA) (fmap f xB) (fmap f xC) (fmap f xD) (fmap f xE) (fmap f xF) 

mergeWith :: Eq a => (b -> c -> d) -> HashMap a b -> HashMap a c -> HashMap a d
mergeWith _ Empty _ = Empty
mergeWith _ _ Empty = Empty
mergeWith f (Bucket xs) (Bucket ys) = emptyOrBucket $ concatMap mrg xs where
  mrg (k,x) = case lookup k ys of
    Just y -> [(k, f x y)]
    Nothing -> [] 
mergeWith f (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF)               
            (Branch y0 y1 y2 y3 y4 y5 y6 y7 y8 y9 yA yB yC yD yE yF) =
  Branch (w x0 y0) (w x1 y1) (w x2 y2) (w x3 y3) (w x4 y4) (w x5 y5) (w x6 y6) (w x7 y7)
         (w x8 y8) (w x9 y9) (w xA yA) (w xB yB) (w xC yC) (w xD yD) (w xE yE) (w xF yF) where
    w x y = mergeWith f x y                                               
          
fromList :: Eq a => [(a,b)] -> HashMap a b         
fromList xs = fold insMap empty xs where
   insMap m (k,v) = insert k v m 

size :: HashMap a b -> Int
size m = length (entries m)      
                
entries :: HashMap a b -> [(a,b)]          
entries Empty = []
entries (Bucket xs) = xs
entries (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF) = 
     concatMap entries [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, xA, xB, xC, xD, xE, xF]
   
keys :: HashMap a b -> [a]    
keys m = map fst $ entries m     
     
instance ListSource HashMap a where
   toList m = map snd $ entries m       
     
private emptyBranch :: HashMap a b
private emptyBranch = Branch Empty Empty Empty Empty Empty Empty Empty Empty
                             Empty Empty Empty Empty Empty Empty Empty Empty

private changeBranch :: HashMap a b -> Int -> (HashMap a b -> HashMap a b) -> HashMap a b
private changeBranch br  0 f = br.{ b0 <- f } 
private changeBranch br  1 f = br.{ b1 <- f } 
private changeBranch br  2 f = br.{ b2 <- f } 
private changeBranch br  3 f = br.{ b3 <- f } 
private changeBranch br  4 f = br.{ b4 <- f } 
private changeBranch br  5 f = br.{ b5 <- f } 
private changeBranch br  6 f = br.{ b6 <- f } 
private changeBranch br  7 f = br.{ b7 <- f } 
private changeBranch br  8 f = br.{ b8 <- f } 
private changeBranch br  9 f = br.{ b9 <- f } 
private changeBranch br 10 f = br.{ bA <- f } 
private changeBranch br 11 f = br.{ bB <- f } 
private changeBranch br 12 f = br.{ bC <- f } 
private changeBranch br 13 f = br.{ bD <- f } 
private changeBranch br 14 f = br.{ bE <- f } 
private changeBranch br  _ f = br.{ bF <- f } 

private getBranch :: HashMap a b -> Int -> HashMap a b
private getBranch (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF) n = case n of 
   0 -> x0; 1 -> x1; 2 -> x2; 3 -> x3; 4 -> x4; 5 -> x5; 6 -> x6; 7 -> x7;
   8 -> x8; 9 -> x9; 10 -> xA; 11 -> xB; 12 -> xC; 13 -> xD; 14 -> xE; _ -> xF

private emptyOrBucket :: [(a,b)] -> HashMap a b
private emptyOrBucket [] = Empty
private emptyOrBucket xs = Bucket xs    